import os
from datetime import datetime

import pandas as pd

from utils.file import get_root_path
from websocket.msghandler import MsgContext, MessageType


class MultipleExcelHandler:
    # 获得项目根目录路径
    excel_data_path = get_root_path() + r'\excel_data'

    @classmethod
    def get_file_path(cls, dir_name, filename=None, suffix=".xlsx"):
        # 检查文件夹是否存在
        if not os.path.exists(dir_name):
            os.makedirs(dir_name)  # 如果文件夹不存在，则创建它
        now = datetime.now()
        if filename is None:
            filename = now.strftime("%Y-%m-%d_%H-%M-%S")
        filename_with_extension = filename + suffix
        return os.path.join(dir_name, filename_with_extension)

    @classmethod
    def handle(cls, message):
        file_path = cls.get_file_path(cls.excel_data_path)
        algorithm_data = {}
        table_value = message['table_value']
        best_value_index = message['best_value_index']
        M = [message['M']] * len(table_value)
        D = [message['D']] * len(table_value)
        statistical_result = message['statistical_result']
        statistical_result[len(statistical_result) - 1][0:0] = ["+/-/=", None, None]
        for index, val in enumerate(message['tableHead']):
            temp = []
            for i in range(len(message['tableRows'])):
                temp.append(table_value[i][index] + statistical_result[i][index])
            algorithm_data[val] = temp
        # 示例数据
        df = pd.DataFrame({
            'Problem': message['tableRows'],
            "M": M,
            "D": D,
            **algorithm_data
        })
        statistical_result[len(statistical_result) - 1].append('')
        df.loc[len(df)] = statistical_result[len(statistical_result) - 1]
        with pd.ExcelWriter(file_path, engine='xlsxwriter') as writer:
            df.to_excel(writer, index=False, sheet_name='Sheet1')
            # 获取 xlsxwriter workbook 和 worksheet 对象
            workbook = writer.book
            cell_format = workbook.add_format({'font_name': 'Times New Roman', 'align': 'center', 'valign': 'vcenter'})
            worksheet = writer.sheets['Sheet1']
            # 设置header不加粗格式
            header_format = workbook.add_format({'bold': False, 'align': 'center', 'valign': 'vcenter'})
            bold_format = workbook.add_format({'bold': True, 'align': 'center', 'valign': 'vcenter'})

            # 遍历表头设置格式（注意：'max_col' 从0开始，故需要+1）
            for col_num in range(df.shape[1]):
                worksheet.write(0, col_num, df.columns[col_num], header_format)
            worksheet.set_column('A:XFD', 25, cell_format=cell_format)
            for row_num in range(0, df.shape[0] - 1):
                for col_num in range(3, df.shape[1]):
                    if best_value_index[row_num][col_num - 3]:
                        worksheet.write(row_num + 1, col_num, df.iat[row_num, col_num], bold_format)
            # 合并单元格操作，合并第4行的第一列到第三列单元格
            worksheet.merge_range(f"A{len(message['tableRows']) + 2}:C{len(message['tableRows']) + 2}", "+/-/=",
                                  workbook.add_format({'align': 'center'}))
        MsgContext.task_complete_return(MessageType.MULTI_DATA_2_EXCEL_COMPLETE, f"file is saved in {file_path}")
